import cv2
import numpy as np
import pyautogui
import pydirectinput
from ultralytics import YOLO
import time
import random
import queue
import threading


class GameNavigator:
    def __init__(self, model_path='best.pt', region=(0, 87, 842, 570),
                 move_cooldown=0.5, move_duration=0.01, alpha_danger=0.5,
                 detection_interval=0.1, post_move_wait=0.1):
        """初始化游戏导航器"""
        self.model = YOLO(model_path)
        self.x, self.y, self.width, self.height = region
        self.move_cooldown = move_cooldown  # 移动冷却时间（秒）
        self.move_duration = move_duration  # 按键持续时间（秒）
        self.alpha_danger = alpha_danger  # 危险区域透明度
        self.detection_interval = detection_interval  # 检测间隔（秒）
        self.post_move_wait = post_move_wait  # 移动后等待时间（秒）

        # 定义移动方向
        self.DIRECTIONS = [(0, -1), (1, 0), (0, 1), (-1, 0)]  # 上，右，下，左
        self.DIRECTION_NAMES = ['上', '右', '下', '左']
        self.current_direction_index = 0

        # 记录上一次移动方向，初始化为None
        self.last_direction_index = None

        # 移动状态
        self.last_move_time = 0
        self.has_moved = False  # 标记是否已移动
        self.move_pending = False  # 标记是否有待执行的移动
        self.move_complete_time = 0  # 上次移动完成的时间

        # 检测结果
        self.current_results = None
        self.detection_queue = queue.Queue(maxsize=1)
        self.last_detection_time = 0

        # FPS计算
        self.fps = 0
        self.last_time = time.time()

        # 启动检测线程
        self.detection_thread = threading.Thread(target=self.detection_loop, daemon=True)
        self.detection_thread.start()

    def detection_loop(self):
        """检测线程，持续处理检测任务"""
        while True:
            try:
                # 限制检测频率
                current_time = time.time()
                if current_time - self.last_detection_time < self.detection_interval:
                    time.sleep(0.01)  # 短暂休眠，避免CPU占用过高
                    continue

                # 从队列获取图像进行检测
                try:
                    image = self.detection_queue.get_nowait()
                    self.current_results = self.model(image)
                    self.detection_queue.task_done()
                    self.last_detection_time = current_time
                except queue.Empty:
                    # 如果队列为空，尝试捕获新图像
                    image = self.capture_screen()
                    if image is not None:
                        self.current_results = self.model(image)
                        self.last_detection_time = current_time

            except Exception as e:
                print(f"检测线程错误: {e}")
                time.sleep(0.1)  # 出错后短暂休眠

    def capture_screen(self):
        """捕获游戏区域屏幕"""
        try:
            screenshot = pyautogui.screenshot(region=(self.x, self.y, self.width, self.height))
            return cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)
        except Exception as e:
            print(f"截图错误: {e}")
            return None

    def process_detection_results(self, results):
        """处理检测结果，提取矿石、岩浆和角色位置"""
        detection_info = {
            '岩浆': [],
            '石块': [],
            '自己': [],
            '障碍物': []  # 新增障碍物类别
        }

        if not results:
            print("未检测到任何对象")
            return [], [], [], []

        for r in results:
            boxes = r.boxes
            for box in boxes:
                class_id = r.names[int(box.cls)]
                confidence = float(box.conf)
                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)

                # 只考虑置信度高于阈值的已知类别
                if confidence > 0.5 and class_id in detection_info:
                    detection_info[class_id].append((x1, y1, x2, y2))
                elif confidence > 0.5:
                    print(f"忽略未知类别: {class_id}")

        return detection_info['石块'], detection_info['岩浆'], detection_info['自己'], detection_info['障碍物']

    def get_role_center(self, role_locations):
        """获取角色中心位置"""
        if not role_locations:
            return None, None

        x1, y1, x2, y2 = role_locations[0]
        center_x = int((x1 + x2) / 2)
        center_y = int((y1 + y2) / 2)
        return center_x, center_y

    def visualize_areas(self, image, ore_locations, magma_locations, obstacle_locations):
        """可视化危险区域和安全区域，并返回危险区域掩码"""
        if image is None:
            print("图像为空，无法可视化")
            return None, None, None

        # 创建颜色掩码和危险区掩码
        color_mask = np.zeros_like(image)
        danger_mask = np.zeros((self.height, self.width), dtype=np.uint8)

        # 获取角色位置（如果已知）
        role_x, role_y = self.get_role_center(self.process_detection_results(self.current_results)[2])

        # 根据角色位置和移动方向动态调整危险区偏移
        if role_x is not None and role_y is not None:
            # 获取当前移动方向
            current_direction = self.DIRECTIONS[self.current_direction_index]

            # 基础偏移值
            base_ore_offset = 0
            base_magma_offset = 0
            base_obstacle_offset = 10  # 障碍物偏移

            # 根据移动方向调整偏移值，确保偏移方向与移动方向一致
            # 修正：添加负号，使偏移方向与移动方向一致
            direction_offset_x = -current_direction[0] * 20  # 横向偏移
            direction_offset_y = -current_direction[1] * 20  # 纵向偏移

            # 角色下方的危险区偏移更大
            vertical_factor = min(1.0, (self.height - role_y) / self.height * 2)

            # 综合计算最终偏移
            ore_offset_x = int(direction_offset_x)
            ore_offset_y = int(base_ore_offset * vertical_factor + direction_offset_y)

            magma_offset_x = int(direction_offset_x * 1.5)  # 岩浆偏移更大
            magma_offset_y = int(base_magma_offset * vertical_factor + direction_offset_y * 1.5)

            obstacle_offset_x = int(direction_offset_x)
            obstacle_offset_y = int(base_obstacle_offset * vertical_factor + direction_offset_y)
        else:
            ore_offset_x = 0
            ore_offset_y = 0
            magma_offset_x = 0
            magma_offset_y = 0
            obstacle_offset_x = 0
            obstacle_offset_y = 0

        # 标记危险区
        for loc in ore_locations:
            x1, y1, x2, y2 = loc
            x1 += ore_offset_x
            x2 += ore_offset_x
            y1 += ore_offset_y
            y2 += ore_offset_y

            # 确保区域在截图内
            if y1 < self.height and y2 > 0 and x1 < self.width and x2 > 0:
                y1 = max(y1, 0)
                y2 = min(y2, self.height)
                x1 = max(x1, 0)
                x2 = min(x2, self.width)
                # 使用矩形代替椭圆
                cv2.rectangle(color_mask, (x1, y1), (x2, y2), (0, 0, 255), -1)
                cv2.rectangle(danger_mask, (x1, y1), (x2, y2), 255, -1)

        for loc in magma_locations:
            x1, y1, x2, y2 = loc
            x1 += magma_offset_x
            x2 += magma_offset_x
            y1 += magma_offset_y
            y2 += magma_offset_y

            if y1 < self.height and y2 > 0 and x1 < self.width and x2 > 0:
                y1 = max(y1, 0)
                y2 = min(y2, self.height)
                x1 = max(x1, 0)
                x2 = min(x2, self.width)
                # 使用矩形代替椭圆
                cv2.rectangle(color_mask, (x1, y1), (x2, y2), (0, 0, 255), -1)
                cv2.rectangle(danger_mask, (x1, y1), (x2, y2), 255, -1)

        # 标记障碍物
        for loc in obstacle_locations:
            x1, y1, x2, y2 = loc
            x1 += obstacle_offset_x
            x2 += obstacle_offset_x
            y1 += obstacle_offset_y
            y2 += obstacle_offset_y

            if y1 < self.height and y2 > 0 and x1 < self.width and x2 > 0:
                y1 = max(y1, 0)
                y2 = min(y2, self.height)
                x1 = max(x1, 0)
                x2 = min(x2, self.width)
                cv2.rectangle(color_mask, (x1, y1), (x2, y2), (255, 0, 0), -1)  # 蓝色表示障碍物
                cv2.rectangle(danger_mask, (x1, y1), (x2, y2), 255, -1)

        # 创建安全区掩码
        safe_mask = np.ones((self.height, self.width), dtype=np.uint8) * 255
        safe_mask[danger_mask > 0] = 0

        # 合并掩码与原图
        result = cv2.addWeighted(image, 1, color_mask, self.alpha_danger, 0)
        return result, safe_mask, danger_mask

    def check_safe_direction(self, safe_mask, current_x, current_y, direction, step_size=10):
        """检查指定方向是否安全"""
        if safe_mask is None or current_x is None or current_y is None:
            return False

        dx, dy = direction
        new_x = current_x + dx * step_size
        new_y = current_y + dy * step_size

        # 确保新位置在安全掩码范围内
        if 0 <= new_x < safe_mask.shape[1] and 0 <= new_y < safe_mask.shape[0]:
            # 检查目标位置周围的区域是否安全
            for i in range(-3, 4):  # 检查目标位置周围7x7的区域
                for j in range(-3, 4):
                    nx = new_x + i
                    ny = new_y + j
                    if 0 <= nx < safe_mask.shape[1] and 0 <= ny < safe_mask.shape[0]:
                        if safe_mask[ny, nx] == 0:
                            return False
            return True
        return False

    def check_direction_danger(self, danger_mask, current_x, current_y, direction, step_size=30, check_size=15):
        """检查指定方向是否存在危险"""
        if danger_mask is None or current_x is None or current_y is None:
            return True  # 无法获取危险掩码时，默认认为有危险

        dx, dy = direction
        new_x = current_x + dx * step_size
        new_y = current_y + dy * step_size

        # 确保新位置在危险掩码范围内
        if not (0 <= new_x < danger_mask.shape[1] and 0 <= new_y < danger_mask.shape[0]):
            return True  # 超出边界，视为危险

        # 计算目标区域内的危险像素数量
        x_min = max(0, new_x - check_size)
        y_min = max(0, new_y - check_size)
        x_max = min(danger_mask.shape[1], new_x + check_size)
        y_max = min(danger_mask.shape[0], new_y + check_size)

        danger_area = danger_mask[y_min:y_max, x_min:x_max]
        danger_pixels = np.sum(danger_area) / 255  # 转换为像素数量

        # 如果危险像素超过阈值，则认为该方向有危险
        return danger_pixels > 10  # 阈值可调整

    def find_safe_adjacent_direction(self, safe_mask, danger_mask, current_x, current_y, current_direction_index):
        """寻找相邻的安全方向，避免选择反方向"""
        if safe_mask is None or danger_mask is None or current_x is None or current_y is None:
            return None, None

        # 获取当前方向的相邻方向（左右）
        left_dir_idx = (current_direction_index - 1) % len(self.DIRECTIONS)
        right_dir_idx = (current_direction_index + 1) % len(self.DIRECTIONS)

        # 计算反方向
        opposite_dir_idx = (current_direction_index + 2) % len(self.DIRECTIONS)

        # 优先检查左侧方向
        if left_dir_idx != opposite_dir_idx:  # 确保不是反方向
            left_direction = self.DIRECTIONS[left_dir_idx]
            if self.check_safe_direction(safe_mask, current_x, current_y, left_direction) and \
                    not self.check_direction_danger(danger_mask, current_x, current_y, left_direction):
                return left_direction, left_dir_idx

        # 再检查右侧方向
        if right_dir_idx != opposite_dir_idx:  # 确保不是反方向
            right_direction = self.DIRECTIONS[right_dir_idx]
            if self.check_safe_direction(safe_mask, current_x, current_y, right_direction) and \
                    not self.check_direction_danger(danger_mask, current_x, current_y, right_direction):
                return right_direction, right_dir_idx

        # 如果左右都不安全或都是反方向，尝试其他方向（排除反方向）
        for i, direction in enumerate(self.DIRECTIONS):
            if i != current_direction_index and i != opposite_dir_idx:
                if self.check_safe_direction(safe_mask, current_x, current_y, direction) and \
                        not self.check_direction_danger(danger_mask, current_x, current_y, direction):
                    return direction, i

        # 如果所有方向都不安全或都是反方向，返回None
        return None, None

    def perform_move(self, direction):
        """执行实际的移动操作"""
        dx, dy = direction
        key_map = {
            (0, -1): 'up',
            (1, 0): 'right',
            (0, 1): 'down',
            (-1, 0): 'left'
        }

        key = key_map.get((dx, dy))
        if key:
            # 记录当前移动方向的索引
            self.last_direction_index = self.DIRECTIONS.index(direction)

            print(f"执行移动: {self.DIRECTION_NAMES[self.last_direction_index]}")

            # 标记移动已执行
            self.has_moved = True
            self.move_pending = False

            # 使用单独的线程执行按键操作，避免阻塞主线程
            threading.Thread(target=self._execute_move, args=(key,), daemon=True).start()

            # 移动后立即触发一次屏幕捕获和检测（在新线程中执行，避免阻塞）
            threading.Thread(target=self._refresh_after_move, daemon=True).start()

    def _execute_move(self, key):
        """在线程中执行实际的按键操作"""
        try:
            # 按下按键
            pydirectinput.keyDown(key)

            # 保持按键一段时间
            time.sleep(self.move_duration)

            # 释放按键
            pydirectinput.keyUp(key)

            # 记录移动完成时间
            self.move_complete_time = time.time()
        except Exception as e:
            print(f"移动执行错误: {e}")

    def calculate_fps(self):
        """计算FPS"""
        current_time = time.time()
        fps = 1.0 / (current_time - self.last_time)
        self.last_time = current_time
        return fps

    def navigate(self):
        """执行导航逻辑"""
        # 计算FPS
        self.fps = self.calculate_fps()

        # 捕获屏幕
        screenshot = self.capture_screen()
        if screenshot is None:
            return

        # 将图像放入检测队列
        try:
            self.detection_queue.put_nowait(screenshot)
        except queue.Full:
            pass  # 队列已满，跳过此次检测

        # 处理检测结果
        if hasattr(self, 'current_results') and self.current_results is not None:
            ore_locations, magma_locations, role_locations, obstacle_locations = self.process_detection_results(
                self.current_results)

            # 获取角色中心位置
            center_x, center_y = self.get_role_center(role_locations)

            # 可视化区域并获取危险区域掩码
            marked_image, safe_mask, danger_mask = self.visualize_areas(screenshot, ore_locations, magma_locations,
                                                                        obstacle_locations)

            # 检查是否需要移动
            current_time = time.time()

            # 如果有待执行的移动，不执行新的决策
            if self.move_pending:
                return

            # 如果已经移动过，检查是否需要强制刷新检测
            if self.has_moved:
                # 移动后立即强制刷新一次检测
                if current_time - self.move_complete_time <= 0.1:  # 移动后0.1秒内
                    print("移动后强制刷新检测")
                    # 捕获新的屏幕并进行检测
                    new_screenshot = self.capture_screen()
                    if new_screenshot is not None:
                        self.current_results = self.model(new_screenshot)
                        ore_locations, magma_locations, role_locations, obstacle_locations = self.process_detection_results(
                            self.current_results)
                        center_x, center_y = self.get_role_center(role_locations)
                        marked_image, safe_mask, danger_mask = self.visualize_areas(new_screenshot, ore_locations,
                                                                                    magma_locations, obstacle_locations)

                # 减少移动后等待时间，允许更快地进行下一步决策
                if current_time - self.move_complete_time >= self.post_move_wait:
                    print(f"移动后等待了 {self.post_move_wait} 秒，准备下一步移动...")
                    self.has_moved = False
                else:
                    print(f"移动后等待中: {current_time - self.move_complete_time:.2f}/{self.post_move_wait}")

            if current_time - self.last_move_time >= self.move_cooldown:
                if center_x is not None and center_y is not None:
                    # 获取当前计划移动的方向
                    current_direction = self.DIRECTIONS[self.current_direction_index]

                    # 检查当前方向是否安全且没有危险
                    is_safe = self.check_safe_direction(safe_mask, center_x, center_y, current_direction)
                    has_danger = self.check_direction_danger(danger_mask, center_x, center_y, current_direction)

                    if is_safe and not has_danger:
                        # 当前方向安全且无危险，计划执行移动
                        self.move_pending = True
                        self.perform_move(current_direction)
                        self.last_move_time = current_time
                        print(f"向 {self.DIRECTION_NAMES[self.current_direction_index]} 移动")
                    else:
                        # 当前方向不安全或有危险，尝试寻找相邻安全方向
                        print(f"当前方向 {self.DIRECTION_NAMES[self.current_direction_index]} 有危险，尝试寻找其他方向")
                        safe_direction, safe_index = self.find_safe_adjacent_direction(
                            safe_mask, danger_mask, center_x, center_y, self.current_direction_index
                        )

                        if safe_direction:
                            # 找到安全方向，计划执行移动
                            self.current_direction_index = safe_index
                            self.move_pending = True
                            self.perform_move(safe_direction)
                            self.last_move_time = current_time
                            print(f"转向 {self.DIRECTION_NAMES[safe_index]} 移动")
                        else:
                            # 没有找到安全方向，随机选择一个方向尝试（避免反方向）
                            available_directions = list(range(len(self.DIRECTIONS)))
                            if self.last_direction_index is not None:
                                # 排除反方向
                                opposite_idx = (self.last_direction_index + 2) % len(self.DIRECTIONS)
                                available_directions = [d for d in available_directions if d != opposite_idx]

                            if available_directions:
                                random_idx = random.choice(available_directions)
                                random_direction = self.DIRECTIONS[random_idx]
                                self.move_pending = True
                                self.perform_move(random_direction)
                                self.current_direction_index = random_idx
                                self.last_move_time = current_time
                                print("没有找到安全方向，随机移动（避免反方向）")
                            else:
                                # 如果没有可用方向（理论上不会发生），仍随机选择
                                random_direction = random.choice(self.DIRECTIONS)
                                self.move_pending = True
                                self.perform_move(random_direction)
                                self.current_direction_index = self.DIRECTIONS.index(random_direction)
                                self.last_move_time = current_time
                                print("没有找到安全方向，随机移动（所有方向被视为危险）")
                else:
                    print("无法确定角色位置，等待下一帧")

            # 显示结果
            if marked_image is not None:
                # 添加FPS信息
                cv2.putText(marked_image, f"FPS: {self.fps:.1f}", (10, 30),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

                # 添加危险程度信息
                if danger_mask is not None:
                    danger_percentage = np.sum(danger_mask) / (255 * self.width * self.height) * 100
                    cv2.putText(marked_image, f"危险程度: {danger_percentage:.1f}%", (10, 60),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

                # 添加移动状态信息
                status = "移动中" if self.move_pending else "思考中"
                cv2.putText(marked_image, f"状态: {status}", (10, 90),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)

                # 添加检测时间信息
                cv2.putText(marked_image, f"检测间隔: {self.detection_interval:.2f}s", (10, 120),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 165, 0), 2)

                # 添加移动后等待信息
                if self.has_moved:
                    wait_progress = f"移动后等待: {current_time - self.move_complete_time:.2f}/{self.post_move_wait}"
                    cv2.putText(marked_image, wait_progress, (10, 150),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 165, 255), 2)

                cv2.imshow('YOLOv8 Inference', marked_image)
                cv2.waitKey(1)

    def _refresh_after_move(self):
        """移动后刷新检测"""
        # 等待一小段时间，确保游戏响应移动
        time.sleep(0.1)

        # 捕获新的屏幕并进行检测
        new_screenshot = self.capture_screen()
        if new_screenshot is not None:
            self.current_results = self.model(new_screenshot)
            print("移动后检测已刷新")


def main():
    """主函数"""
    print("初始化游戏导航器...")
    # 可以调整detection_interval和post_move_wait参数来优化性能
    navigator = GameNavigator(detection_interval=0.05, post_move_wait=0.6)

    print("按 'q' 键退出程序")
    print("开始导航...")

    try:
        while True:
            navigator.navigate()

            # 检查是否按下 'q' 键退出
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

    except KeyboardInterrupt:
        print("检测到键盘中断，程序退出")
    finally:
        # 确保所有按键都已释放
        for key in ['up', 'down', 'left', 'right']:
            try:
                pydirectinput.keyUp(key)
            except:
                pass

        cv2.destroyAllWindows()
        print("程序已安全退出")


if __name__ == "__main__":
    main()